% Opspace implemenation with DCA
clear

% All joints are revolute in this version
% The first body is attached to ground, the last body is free

% Fixed parameters for the system. Each body in the system is identical
NBODIES = 10;
MASS = 10;
INERTIA = 10;

T = (1:NBODIES)';
T = T+20;

% Center of mass located at the center of each link.
LENGTH = 1;

GRAVITY = [0 -9.81 0]';

% Initialize the state vector of the system
q(1:NBODIES,1) = 1;
u(1:NBODIES,1) = 1;

% For a revolute joint about positive z axis
% H is referred to as P^J_k in our paper. 
H = [0 0 1 0 0 0]';

% D = D^J_k
D = [1 0 0 0 0;
     0 1 0 0 0;
     0 0 0 0 0;
     0 0 1 0 0;
     0 0 0 1 0;
     0 0 0 0 1];
 
Dk = [0 0 ;
      0 0;
      0 0 ;
      1 0 ;
      0 1 ;
      0 0 ];

% PREPROCESSING loop. Operations in this loop are performed just once and
% need not be timed.
for i=1:NBODIES
    
    % Initialize the fixed parameters of the system
    mass(i) = MASS;
    inertia(:,:,i) = [1 0 0;
                      0 1 0;
                      0 0 INERTIA];
                  
    % Spatial mass matrix written in body basis about the center of mass
    M(:,:,i) = [inertia(:,:,i) zeros(3,3);
                zeros(3,3) mass(i)*eye(3,3)];
    
    length(i) = LENGTH;
    
    % Vectors locating handles 1 and 2 with respect to center of mass
    % Written in body basis. 
    cmtoH1(:,i) = [-length(i)/2 0 0]';
    cmtoH2(:,i) = [length(i)/2 0 0]';
    
    % Written in body basis
    cmtoH1cross(:,:,i) = crossmat(cmtoH1(:,i)); 
                     
    cmtoH2cross(:,:,i) = crossmat(cmtoH2(:,i));
                     
    % Shift matrices for spatial constraint forces on handles H1 and H2
    SfH1(:,:,i) = [eye(3) cmtoH1cross(:,:,i);zeros(3) eye(3)];
    SfH2(:,:,i) = [eye(3) cmtoH2cross(:,:,i);zeros(3) eye(3)];
    
    % handle equations for each body are 
    % A_1 = zeta11 F1 + zeta12 F2 + zeta13
    % A_2 = zeta21 F1 + zeta22 F2 + zeta23
    % phi11, phi12, phi21 and phi22 are constant in body basis and computed
    % just once as a preprocessing step. 
    
    shift_inverse_H1(:,:,i) = (SfH1(:,:,i)')*inv(M(:,:,i));
    shift_inverse_H2(:,:,i) = (SfH2(:,:,i)')*inv(M(:,:,i));
    
    zeta11(:,:,i) = shift_inverse_H1(:,:,i)*SfH1(:,:,i); % (1a)
    zeta12(:,:,i) = shift_inverse_H1(:,:,i)*SfH2(:,:,i); % (1a)    
    zeta14(:,i) = shift_inverse_H1(:,:,i)*H;
    zeta15(:,i) = shift_inverse_H1(:,:,i)*H;
    Tzeta14(:,:,i) = [zeta14(:,i) zeta15(:,i)];
    
    zeta21(:,:,i) = shift_inverse_H2(:,:,i)*SfH1(:,:,i); % (1b)
    zeta22(:,:,i) = shift_inverse_H2(:,:,i)*SfH2(:,:,i); % (1b)
    zeta24(:,i) = shift_inverse_H2(:,:,i)*H;
    zeta25(:,i) = shift_inverse_H2(:,:,i)*H;
    Tzeta24(:,:,i) = [zeta24(:,i) zeta25(:,i)];
end

for i=1:NBODIES
    % Transformation matrix between body and parent
    pr_c_k(:,:,i) = [cos(q(i)) -sin(q(i)) 0;
                     sin(q(i)) cos(q(i)) 0;
                     0         0         1];
                 
    % Form spatial trasnformation matrices
    Pr_C_K(:,:,i) = [pr_c_k(:,:,i) zeros(3,3);
                    zeros(3,3) pr_c_k(:,:,i)];
    
    % Transformation matrix between body and newtonian frame
    n_c_k(:,:,i) = [cos(sum(q(1:i))) -sin(sum(q(1:i))) 0;
                     sin(sum(q(1:i))) cos(sum(q(1:i))) 0;
                     0                0                1];
    N_C_K(:,:,i) = [n_c_k(:,:,i) zeros(3,3);
                    zeros(3,3)   n_c_k(:,:,i)];
    % Angular velocity of each body is trivially obtained here
    W(:,i) = [0 0 sum(u(1:i))]';
    
    Fa  = [ -cross(W(:,i),inertia(:,:,i)*W(:,i)); % -w x Iw term
                  mass(i)*n_c_k(:,:,i)'*GRAVITY];% Only external force is gravity
              
    At_H1  = [zeros(3,1);cross(W(:,i),cross(W(:,i),cmtoH1(:,i)))];
    At_H2  = [zeros(3,1);cross(W(:,i),cross(W(:,i),cmtoH2(:,i)))];
    
    % phi13 and phi23 are also written in body basis
    zeta13(:,i) = shift_inverse_H1(:,:,i)*Fa + At_H1; % (1a)
    zeta23(:,i) = shift_inverse_H2(:,:,i)*Fa + At_H2; % (1b)    
end


% A convenient tmp variable.
zeta11_asm_tmp = zeta11(:,:,NBODIES);
zeta12_asm_tmp = zeta12(:,:,NBODIES);
zeta13_asm_tmp = zeta13(:,NBODIES);
zeta14_asm_tmp = Tzeta14(:,1,NBODIES);

zeta21_asm_tmp = zeta21(:,:,NBODIES);
zeta22_asm_tmp = zeta22(:,:,NBODIES);
zeta23_asm_tmp = zeta23(:,NBODIES);
zeta24_asm_tmp = Tzeta24(:,1,NBODIES);

% Assembly
% NBODIES are connected by (NBODIES-1) joint excluding the one which connects
% the first body to ground (if at all)
for i=NBODIES:-1:2   
        
    % Convert handle equations of body i to basis of (i-1)
    % Handle 1
    zeta11_c = Pr_C_K(:,:,i)*zeta11_asm_tmp*Pr_C_K(:,:,i)';
    zeta12_c = Pr_C_K(:,:,i)*zeta12_asm_tmp*Pr_C_K(:,:,i)';
    zeta13_c = Pr_C_K(:,:,i)*zeta13_asm_tmp;
    zeta14_c = Pr_C_K(:,:,i)*zeta14_asm_tmp;    
    
    % Handle 2
    zeta21_c = Pr_C_K(:,:,i)*zeta21_asm_tmp*Pr_C_K(:,:,i)';
    zeta22_c = Pr_C_K(:,:,i)*zeta22_asm_tmp*Pr_C_K(:,:,i)';
    zeta23_c = Pr_C_K(:,:,i)*zeta23_asm_tmp;               
    zeta24_c = Pr_C_K(:,:,i)*zeta24_asm_tmp;    
    
    % Assemble bodies i and (i-1)   
    X_tilde = inv(D'*(zeta11_c + zeta22(:,:,i-1))*D); % (3e)
    X = D*X_tilde*D'; % (3d)
    
    Wa(:,:,i) = X*zeta21(:,:,i-1);        % (3b)
    Wb(:,:,i) = -X*zeta12_c;               % (3b)
    Wc(:,i) = X*(zeta23(:,i-1)-zeta13_c); % (3c)
    
    Wd{i} = X*[zeta24(:,i-1) -(zeta25(:,i-1)+zeta14_c(:,1)) -zeta14_c(:,2:end)];
    
    % Note in equation (3c), \dot{P}^J_K u is identically equal to zero for
    % all planar revolute joints and not included in calculations.
    % PS: The sign of equation Wb in (3b) in the paper draft i sent you  
    % before was wrong.
    
    % Below are equations (4a)-(4f)
    zeta11_asm(:,:,i) = zeta11(:,:,i-1) - zeta12(:,:,i-1)*Wa(:,:,i);
    zeta12_asm(:,:,i) = -zeta12(:,:,i-1)*Wb(:,:,i);
    zeta13_asm(:,i) = zeta13(:,i-1) - zeta12(:,:,i-1)*Wc(:,i) ;
    zeta14_asm{i} = -zeta12(:,:,i-1)*Wd{i};
    zeta14_asm{i}(:,NBODIES-i+1) = zeta14_asm{i}(:,NBODIES-i+1) + zeta14(:,i-1);
    zeta14_asm{i}(:,NBODIES-i+2) = zeta14_asm{i}(:,NBODIES-i+2) - zeta15(:,i-1);
    
    zeta21_asm(:,:,i) = zeta21_c*Wa(:,:,i);             
    zeta22_asm(:,:,i) = zeta22_c + zeta21_c*Wb(:,:,i);  
    zeta23_asm(:,i) = zeta23_c + zeta21_c*Wc(:,i);      
    zeta24_asm{i} = zeta21_c*Wd{i};
    zeta24_asm{i}(:,end-NBODIES+i:end) = zeta24_asm{i}(:,end-NBODIES+i:end) + zeta24_c;        
    
    zeta11_asm_tmp = zeta11_asm(:,:,i);
    zeta12_asm_tmp = zeta12_asm(:,:,i);
    zeta13_asm_tmp = zeta13_asm(:,i);
    zeta14_asm_tmp = zeta14_asm{i};
    
    zeta21_asm_tmp = zeta21_asm(:,:,i);
    zeta22_asm_tmp = zeta22_asm(:,:,i);
    zeta23_asm_tmp = zeta23_asm(:,i);  
    zeta24_asm_tmp = zeta24_asm{i};
    
end

XX = -D*inv(D'*zeta11_asm(:,:,2)*D)*D';
psi1 = XX*zeta13_asm(:,2);
psi2 = XX*zeta14_asm{2};
psi3 = zeta21_asm(:,:,2)*psi1 + zeta23_asm(:,2);
psi4 = zeta21_asm(:,:,2)*psi2 + zeta24_asm{2};
psi5 = Dk'*psi3;
psi6 = Dk'*psi4;


% T = -psi6'*inv(psi6*psi6')*psi5;
% Tn = null(psi6);
% T = T + 100*Tn(:,1);



% Fh1 = psi1 + psi2*T;
% Ah1 = zeta11_asm(:,:,2)*Fh1 + zeta13_asm(:,2) + zeta14_asm{2}*T;
% Ah2 = zeta21_asm(:,:,2)*Fh1 + zeta23_asm(:,2) + zeta24_asm{2}*T;

% Spatial acceleration of ground/parent
Ah2_parent = zeros(6,1);

% Disassembly
for i=1:NBODIES
    
    Ah2_parent_inchildbasis = Pr_C_K(:,:,i)'*Ah2_parent;
    
    if(i~=NBODIES)
        % For all intermediate assembled systems, the spatial constraint
        % force on handle 2 is zero. Hence the equation of handle 1 of an
        % assembled intermediate system can be written as 
        % Ah1 = zeta11_asm*Fh1 + zeta13
        % Since D'*(Ah1 - Ah2_parent) = 0 the above equation reduces to
        % D'*(zeta11_asm*Fh1 + zeta13-Ah2_parent)=0 From this, Fh1 can be
        % calculated as 
        Fh1 = D*((D'*zeta11_asm(:,:,i+1)*D) \ (-D'*(zeta13_asm(:,i+1)+zeta14_asm{i+1}*T(i:end,1)-Ah2_parent_inchildbasis)));
        
        % the stored values of Wa,Wb and Wc are used to calculate the
        % constrained force on hte connecting joint.
        Fh2 = -Wa(:,:,i+1)*Fh1 -Wc(:,i+1) - Wd{i+1}*T(i:end,1); % This is the negative of eq (3a)        
        Ah1 = zeta11_asm(:,:,i+1)*Fh1 + zeta13_asm(:,i+1) + zeta14_asm{i+1}*T(i:end,1);
        
        Ah2 = zeta21(:,:,i)*Fh1 + zeta22(:,:,i)*Fh2 + zeta23(:,i) + zeta24(:,i)*T(i,1) - zeta25(:,i)*T(i+1,1);
        Ah2_parent = Ah2;
    else
        % For the last body, there is no assembly.
        Fh1 = D*((D'*zeta11(:,:,i)*D) \ (-D'*(zeta13(:,i)+zeta14(:,i)*T(end,1)-Ah2_parent_inchildbasis)));
        Ah1 = zeta11(:,:,i)*Fh1 + zeta13(:,i) + zeta14(:,i)*T(end,1);
        Ah2 = zeta21(:,:,i)*Fh1 + zeta23(:,i) + zeta24(:,i)*T(end,1);
    end
    

    % w_parent x w_child terms is identically equal to zero and not taken
    % into account here
    udot(i,1) = H'*(Ah1 - Ah2_parent_inchildbasis);            
end